from mongoengine import *
import pandas as pd
import json

connect('htstool')


####################
# low level classes

class Control(Document):
	POS='positive'
	NT='negative'
	NON='non targeting'
	UN='untreated'
	OTH='other'

	CTRL_CHOICES = (
		(POS, 'Positive controls'),
		(NT, 'Negative controls'),
		(NON, 'Non targeting controls'),
		(UN, 'Untreated controls'),
		(OTH, 'Other controls'),
		)
	# letter of position
	letter=StringField()
	# row number of position 
	num=StringField()
	substancename=StringField()
	
	def position(self):
		return self.letter+self.num
	pos=property(position)	

	meta = {'allow_inheritance': True}

class Dataset(Document):
	'''
	generic grid data, in shape of the plate, 
	can contain numeric as well as string values
	'''
	data=DictField()	

	def to_df(self):
		return pd.DataFrame.from_dict(self.data)
	dataframe=property(to_df)
	
	def to_html(self):
		return pd.DataFrame.from_dict(self.data).to_html()
	table=property(to_html)

	def to_json(self):
		return pd.DataFrame.from_dict(self.data).to_json(orient='index')
	json=property(to_json)

	def to_json_objs(self):
		platedf = self.data
		objs=[]
		for col in platedf.columns:
			for row in platedf.index:
				objs.append({'val':platedf.loc[row, col], 'row':row, 'col':col})
		return json.dumps(objs)
	jsonobjs=property(to_json_objs)

class Plate(Document):
	platename=StringField()
	cellline=StringField()
	condition=StringField()
	concentration=StringField()
	dimension=ListField(IntField())
	lib=StringField()
	
	replicate=IntField()

	###### warning NULLIFY seems to work when deleting Plate object #####
	#controls=ListField(ReferenceField(Control, reverse_delete_rule=NULLIFY))
	controls=ListField(ReferenceField(Control))
	
	data=ReferenceField(Dataset, reverse_delete_rule=CASCADE)

	annotation=ReferenceField(Dataset, reverse_delete_rule=CASCADE)

	def getNTcontrols(self):
		nt_ctrls=filter(lambda ctrl:isinstance(ctrl, NegativeControl), self.controls)
		return nt_ctrls
	
	def getUntreatedcontrols(self):
		nt_ctrls=filter(lambda ctrl:isinstance(ctrl, UntreatedControl), self.controls)
		return nt_ctrls

	def getAllcontrols(self):
		ctrls={}
		for ctrl in self.controls:
			if isinstance(ctrl, PositiveControl):
				ctrls[ctrl.letter+ctrl.num]={'type':'pos', 'desc':ctrl.effect}
			if isinstance(ctrl, NegativeControl):
				ctrls[ctrl.letter+ctrl.num]={'type':'nt', 'desc':ctrl.effect}
			if isinstance(ctrl, NonTargetControl):
				ctrls[ctrl.letter+ctrl.num]={'type':'nontarg', 'desc':ctrl.description}
			if isinstance(ctrl, UntreatedControl):
				ctrls[ctrl.letter+ctrl.num]={'type':'untreated', 'desc':ctrl.description}
			if isinstance(ctrl, OtherControl):
				ctrls[ctrl.letter+ctrl.num]={'type':'other', 'desc':ctrl.description}
		return json.dumps(ctrls)
	meta = {'allow_inheritance': True}

####################
# inherited classes

class PositiveControl(Control):
	concentration=IntField()
	# e.g rad sensitiser
	effect=StringField()

class NegativeControl(Control):
	effect=StringField()

class NonTargetControl(Control):
	description=StringField()

class UntreatedControl(Control):
	description=StringField()

class OtherControl(Control):
	description=StringField()




class RawPlate(Plate):
	measure=StringField()
	channel=StringField()

#class RawPlateMulti(Plate):
#	pass

# depending on how the stat is calculated the process can be either:
# raw -> norm -> summ or
# raw -> summ -> norm
class NormPlate(Plate):
	PERPLATE='perplate-z'
	ROBUSTZ='robust-z'
	BSCORE='bscore'
	PERPLATE_NT='perplate-nt-z'
	PERPLATE_UNTREATED='perplate-untreated-z'
	GLOBAL_NT='global-nt-z'

	NORM_CHOICES = (
		(PERPLATE, 'Standard Z score per-plate'),
		(PERPLATE_NT, 'Z score normalise to NT per-plate'),
		(PERPLATE_UNTREATED, 'Z score normalise to untreated per-plate'),
		(GLOBAL_NT, 'Z score normalised to global NT'),
		(BSCORE, 'B-score normalisation'),
		(ROBUSTZ, 'Robust Z score normalisation'),
			)

	normmethod=StringField(choices=NORM_CHOICES)
	plate_normalised=ReferenceField(Plate)

class SummPlate(Plate):
	"""
	summarisation methods, typically drug plate vs non drug plate
	"""
	SUBTRACT='subtract'
	RATIO='ratio'

	SUMM_CHOICES = (
		(SUBTRACT, 'drug - nondrug'),
		(RATIO,    'drug / nondrug'),
			)


	summmethod=StringField(choices=SUMM_CHOICES)
	plates_summed=ListField(ReferenceField(Plate))


class RankProdPlate(Plate):
	"""
	virtual plate that contains multi datasets from RankProd anaylsis
	"""
	RP_CHOICES = (
		('rankprod', 'rank product'),
			)

	# this is a list of plates of different replicates
	# TODO - may need to use no_dereference in views
	replicate=ListField(ReferenceField(Plate))

	# TODO - add choices for ranking_rule
	# get info from plates being used for rankprod
	ranking_rule=StringField()

	rankprod=ReferenceField(Dataset, reverse_delete_rule=NULLIFY)
	pfp=ReferenceField(Dataset, reverse_delete_rule=NULLIFY)
	pval=ReferenceField(Dataset, reverse_delete_rule=NULLIFY)


class RPSuperPlateSummary(RankProdPlate):
	"""
	virtual summary dataset for rank prod merged plates result.
	data arranged as rows: compounds, cols: replicates
	"""
	plateids=ListField(StringField(max_length=30))
	ranked_ind=ListField(StringField(max_length=30))

	def getAnno(self):
		if len(self.ranked_ind)>0:
			return json.dumps(self.annotation.dataframe.loc[self.ranked_ind]['Annotation'].tolist())
		else:
			return json.dumps([])
	def getTopHits(self, topN=100):
		rp_colname = self.rankprod.dataframe.columns[0]
		result_table = pd.concat([self.data.dataframe, 
			   self.rankprod.dataframe,
			   self.pfp.dataframe,
			   self.pval.dataframe], axis=1)
		#n=len(result_table)/len(self.plateids)
		#plateids_index = reduce(lambda x,y:x+y, [[i]*n for i in self.plateids])
		plateid_index = [label.split('-')[1] for label in self.data.dataframe.index]
		result_table.insert(len(result_table.columns), 'plateid', plateid_index)
		res_df = result_table.sort([rp_colname])[:topN]
		self.ranked_ind=res_df.index.tolist()
		self.save()
		return res_df
	
	# output html of the result table 
	def tbl_result(self):
		return self.getTopHits().to_html()

	res_table=property(tbl_result)
	anno_json=property(getAnno)

###################
# plate container classes
	
class PlateGroup(Document):
	RAW='raw'
	NORM='norm'
	SUMM='summ'
	ANALYSIS='analysis'
	PLATE_TYPES = (
		(RAW, 'raw readout plates'),	
		(NORM, 'normalised plates'),	
		(SUMM, 'summarised plates'),	
		(ANALYSIS, 'results from analysis plates'),	
			)

	groupname=StringField()
	groupdesc=StringField()

	plate_type=StringField(choices=PLATE_TYPES)
	# The common attr that these plates share, e.g condition, cellline etc
	# the dict contains the attr name and the string such as {'condition':'nonrad'}
	platecommonattr=DictField()

	plates=ListField(ReferenceField(Plate, reverse_delete_rule=PULL))
	
	def delete_all(self):
		dataset_ids=[]
		if self.plate_type==PlateGroup.ANALYSIS and isinstance(self.plates[0], RankProdPlate):	
			dataset_ids = dataset_ids+[str(p.rankprod.id) for p in self.plates]
			dataset_ids = dataset_ids+[str(p.pfp.id) for p in self.plates]
			dataset_ids = dataset_ids+[str(p.pval.id) for p in self.plates]
		else:
			dataset_ids = [str(p.data.id) if p.data is not None else None for p in self.plates]

		plate_ids = [str(p.id) for p in self.plates]
		
		if dataset_ids[0] is not None:
			Dataset.objects(id__in=dataset_ids).delete()
		Plate.objects(id__in=plate_ids).delete()


class PlateSetting(Document):
	user = StringField()
	treatment = StringField()
	condition = StringField()
	ctrls = StringField()


class Experiment(Document):
	
	expname=StringField()
	group=StringField()
	screen=StringField()
	screentype=StringField()
	library=StringField()
	description=StringField()

	plategroups=ListField(ReferenceField(PlateGroup, reverse_delete_rule=PULL))

	platesetting=ReferenceField(PlateSetting)




###################
# special classes

#class ReplicateValueArray(EmbeddedDocument):
#	well_label=StringField()
#	plateid=StringField()
#	values=ListField()
#	anno=StringField()
#
#class CombinedPlate(Document):
#	plateids=ListField()
#	value_arrays=ListField(EmbeddedDocumentField(ReplicateValueArray))
#	description=StringField()
#





